import prettytable as pt
import torch
from utils.timer import Timer
import torch.nn as nn
import transformers
from sklearn.metrics import precision_recall_fscore_support, f1_score
from torch.utils.data import DataLoader
import model_utils.build_dataset as BuildDataset
from model_utils import build_dataset
from utils.set_seed import set_seed
import os
from utils.logger import Logger
from model_utils.build_model import AssembleModel
from utils.get_parser import get_basic_parser_for_inference
os.environ["CUDA_DEVICE_ORDER"] = 'PCI_BUS_ID'
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

class Predictor(object):
    def __init__(self, model):
        self.device = torch.device('cuda:0')
        self.model = model.to(self.device)

    def predict(self, epoch, data_loader, data):
        self.model.eval()
        result = []
        i = 0
        with torch.no_grad():
            for data_batch in data_loader:
                pass        
        
    def load(self, path):  # 加载权重
        weight = torch.load(path, map_location=self.device)
        self.model.load_state_dict({k.replace('module.',''):v for k,v in weight.items()})

def prepare_dataset(config, f_path):
        dataset, ori_data = BuildDataset(config)
        _iter = DataLoader(dataset=dataset,
                           batch_size=config.batch_size,  
                           drop_last='train' in 'f_path',
                           shuffle=False,
                           num_workers=config.workers, 
                           collate_fn=build_dataset.collate_fn) 
        return _iter, len(dataset), ori_data


if __name__ == '__main__':
    parser = get_basic_parser_for_inference()

    config = parser.parse_args()
    timer = Timer(config.bert_name)
    print("Loading Data")
    test_loader, _total,ori_data=prepare_dataset(config, config.test_dir) 
    table = pt.PrettyTable(['data_dir', 'number'])
    table.add_row([config.test_dir, _total])
    print(table)
    model = AssembleModel(config)
    predictor = Predictor(model)
    predictor.load(config.save_path)
    print('---------Initial model Done.---------')
    predictor.predict("Final", test_loader, ori_data)
    timer.finish()